#!/usr/bin/python3
import math
import copy
import numpy as np
from pathlib import Path



class VectorSpaceModel:

    # constant
    prefix_docid = 'docid_'
    top_k = 10
    threshold_score = 0.01
    num_ret_least = 2

    # docs的模型内部存储格式
    docs = {}
    # 正向索引
    idx = {} 
    # 倒排索引, 需要预先定义好的vocabulary, 或自建vocabulary
    idx_i = {} 

    def __init__(self, docs):
        self.init(docs)

    def __str__(self):
        return str(self.idx)

    def init(self, docs):
        if type(docs) != dict:
            raise ValueError('supplied argument should be of type dict')
        docs_original = {}
        docs_original = copy.deepcopy(docs)

        self.docs = self.construct_docs(docs_original)
        self.idx = self.construct_idx(self.docs)
        self.idx_i = self.construct_inverse_idx(self.docs)
        print('index construct successfully')

    # 更新索引：文档库中增加新文档
    def update(self, doc):
        if type(doc) != str:
            raise ValueError('supplied argument should be of type str')
        # TODO 

        return True
    
    # user query: return docs-like dict
    def query(self, query):
        if type(query) != str:
            raise ValueError('supplied argument should be of type str')
        
        rst = {}
        scores = {}
        '''
        {
            'docid0': 10,
            'docid1': 20
            ...
        }
        '''
        
        # score the relevant docs
        con_q = self.concordance(query)
        terms_q = con_q.keys()

        # 正向索引查找评分法
        # for doc_id, con_d in self.idx.items():
        #     s = self.score(con_q, con_d)
        #     scores[doc_id] = s
        

        # 反向索引查找评分
        for term in terms_q:
            tf_q = con_q.get(term)

            if term in self.idx_i: # 此查询term 有doc包含 计算分数
                df = self.idx_i[term]['df']  # 利用到idf时使用

                l_docs = self.idx_i[term]['docs']
                for doc in l_docs:
                    docid = doc['docid']
                    if docid in scores:
                        scores[docid] += tf_q * doc['weight']
                    else:
                        scores[docid] = tf_q * doc['weight']
        # normalization
        q_norm = self.magnitude(con_q)
        for docid, score in scores.items():
            con_d = self.idx[docid]
            d_norm = self.magnitude(con_d)
            scores[docid] /= (q_norm * d_norm)
        

        # rank
        ranked_scores_list = self.rank(scores)
        # filter
        filtered_scores_list = self.filter(ranked_scores_list)
        # enough rst? if not, return the list before filtered
        docslist = filtered_scores_list if (len(filtered_scores_list) >= self.num_ret_least) else ranked_scores_list[:self.num_ret_least]
        
        # make final return
        for doc_id, score in docslist:
            rst[doc_id] = {}
            rst[doc_id]['score'] = score
            rst[doc_id]['doc_content'] = self.docs.get(doc_id)
         
        return rst

    # 自建内部docs库，方便自定义格式和修改
    def construct_docs(self, docs):
        if type(docs) != dict:
            raise ValueError('supplied argument should be of type dict')

        docs_rst = {}
        for k, v in docs.items():
            # some modification on docs format and content
            doc_id = str(k) if str(k).startswith(self.prefix_docid) else (self.prefix_docid + str(k))
            doc_content = v

            docs_rst[doc_id] = doc_content
        return docs_rst

    def construct_idx(self, docs):
        if type(docs) != dict:
            raise ValueError('supplied argument should be of type dict')
        
        idx = {}
        for k, v in docs.items():

            doc_id = str(k) if str(k).startswith(self.prefix_docid) else (self.prefix_docid + str(k))
            con = self.concordance(v)
            idx[doc_id] = con

        return idx
    
    # 倒排索引：以term为键，行存储，存储doc列表，每个doc项为一个dict，包含相关weight
    '''
    {
        't1': {
            'df':100,
            'docs':[
                {
                    'docid':'doc_...',
                    'termid':'t1',
                    'weight':2,
                    'other_t-d_feature':'...'
                },
                {

                }
            ],
            'other_term_feature':'...'
        },
        't2':{

        }
    }
    '''
    def construct_inverse_idx(self, docs: dict):
        if type(docs) != dict:
            raise ValueError('supplied argument should be of type dict')
        idx_i = {}
        for docid, doc in docs.items():
            con = self.concordance(doc)
            for word, weight in con.items():
                word = word.lower()
                # make a doc_node
                doc_dict = {}
                doc_dict['docid'] = docid
                doc_dict['termid'] = word
                doc_dict['weight'] = weight

                if word in idx_i: # append the doc_node into this term
                    term_info = idx_i[word]
                    term_info['df'] += 1
                    docs_list = term_info['docs']
                    docs_list.append(doc_dict)
                else: # add this word as a term into the index
                    term_info = {}
                    term_info['df'] = 1
                    docs_list = []
                    docs_list.append(doc_dict)
                    term_info['docs'] = docs_list
                    # 将此term节点加入反向索引
                    idx_i[word] = term_info
        # print(idx_i)         
        return idx_i

    # 评分结果排序函数
    def rank(self, scores_dict):
        if type(scores_dict) != dict:
            raise ValueError('supplied argument should be of type dict')
        # reture a sorted list
        l_k = sorted(scores_dict.items(), key=lambda item: item[0], reverse=True)
        l_v = sorted(scores_dict.items(), key=lambda item: item[1], reverse=True)
        return l_v

    # 过滤函数
    def filter(self, l):
        rst = []
        # top k
        rst = l[:self.top_k]
        # threshold_score 
        rst = [(k, v) for (k, v) in rst if v > self.threshold_score]

        return rst


    '''
    len of a vector
    '''
    def magnitude(self, concordance):
        if type(concordance) != dict:
            raise ValueError('supplied argument should be of type dict')
        
        total = 0
        for word, weight in concordance.items():
            total += weight**2
        return math.sqrt(total)


    '''
    score-function: based on relation function
        when score, query is treated as a short document to calculate the relevance based on relation-function,
        but this is just a hypothesis, if we want to make some adjustment on score-calculation
        say, do some preprocessing on query vector for different application, we should not directly modify the relation-func
        so we can abstract a function, named score(q), to do this
    
    '''
    def score(self, query, doc):
        if type(query) != dict:
            raise ValueError('supplied argument should be of type dict')
        if type(doc) != dict:
            raise ValueError('supplied argument should be of type dict')

        # do some preprocessing
        q = self.preprocess(query)
        s = self.relation(q, doc)
        return s

    # 查询向量预处理
    def preprocess(self, query):

        q = query

        return q

    '''
    doc relation calculation(no specific processing): cosine-similarity (vector dot product)
    '''
    def relation(self, concordance_1, concordance_2):
        if type(concordance_1) != dict:
            raise ValueError('supplied argument should be of type dict')
        if type(concordance_2) != dict:
            raise ValueError('supplied argument should be of type dict')

        dot_product = 0.0
        for word, weight in concordance_1.items():
            if word in concordance_2:
                dot_product += weight * concordance_2[word]

        normalization = self.magnitude(concordance_1) * self.magnitude(concordance_2)

        if normalization != 0:
            return dot_product / normalization
        else:
            return 0

    '''
    vector of doc : weight is tf
    '''
    def concordance(self, document):
        if type(document) != str:
            raise ValueError('supplied argument should be of type string')
        
        con = {}
        for word in document.split(' '):
            word = word.lower() # index word all to lower format

            if word in con:
                con[word] += 1
            else:
                con[word] = 1
        return con

# test doc
test_doc1 = 'this is my search engine and recommendation homework , enjoy it, peace'
test_doc2 = 'it is just a rudimentary toy program, hope it works, peace'
test_query1 = 'search engine peace'
test_query2 = 'CAPTCHA'

documents = {}
documents_test = {
  0:'''At Scale You Will Hit Every Performance Issue I used to think I knew a bit about performance scalability and how to keep things trucking when you hit large amounts of data Truth is I know diddly squat on the subject since the most I have ever done is read about how its done To understand how I came about realising this you need some background''',
  1:'''Richard Stallman to visit Australia Im not usually one to promote events and the like unless I feel there is a genuine benefit to be had by attending but this is one stands out Richard M Stallman the guru of Free Software is coming Down Under to hold a talk You can read about him here Open Source Celebrity to visit Australia''',
  2:'''MySQL Backups Done Easily One thing that comes up a lot on sites like Stackoverflow and the like is how to backup MySQL databases The first answer is usually use mysqldump This is all fine and good till you start to want to dump multiple databases You can do this all in one like using the all databases option however this makes restoring a single database an issue since you have to parse out the parts you want which can be a pain''',
  3:'''Why You Shouldnt roll your own CAPTCHA At a TechEd I attended a few years ago I was watching a presentation about Security presented by Rocky Heckman read his blog its quite good In it he was talking about security algorithms The part that really stuck with me went like this''',
  4:'''The Great Benefit of Test Driven Development Nobody Talks About The feeling of productivity because you are writing lots of code Think about that for a moment Ask any developer who wants to develop why they became a developer One of the first things that comes up is I enjoy writing code This is one of the things that I personally enjoy doing Writing code any code especially when its solving my current problem makes me feel productive It makes me feel like Im getting somewhere Its empowering''',
  5:'''Setting up GIT to use a Subversion SVN style workflow Moving from Subversion SVN to GIT can be a little confusing at first I think the biggest thing I noticed was that GIT doesnt have a specific workflow you have to pick your own Personally I wanted to stick to my Subversion like work-flow with a central server which all my machines would pull and push too Since it took a while to set up I thought I would throw up a blog post on how to do it''',
  6:'''Why CAPTCHA Never Use Numbers 0 1 5 7 Interestingly this sort of question pops up a lot in my referring search term stats Why CAPTCHAs never use the numbers 0 1 5 7 Its a relativity simple question with a reasonably simple answer Its because each of the above numbers are easy to confuse with a letter See the below''',
  7:'''Why CAPTCHA CAPTCHAs a reasonably simple answer Its because each of the above numbers are easy to confuse with a letter See the below''',
}

# data : E:\资料\aclImdb_v1\aclImdb\train\unsup 前100条数据
def prepare_docs():

    global documents
    root = Path('./')
    data_path = root / 'data'
    i = 0
    for p in data_path.iterdir():
        doc = p.read_text(encoding='utf8')
        documents[i] = doc
        i += 1
    # print(documents)
    print('prepared data successfully')

def main():
    prepare_docs()
    model = VectorSpaceModel(documents)
    # l = list(model.idx_i.items())
    # print(l[:2])
    q = ''
    STOP = '0'
    while q != STOP:
        q = input('\nplease input your query:')
        rst = model.query(q)
        print(rst)


if __name__ == "__main__":
    main()

